import pandas as pd
import pulp as pp
from tqdm import tqdm

data_path = 'sc60.csv'

try:
    hourly_prediction = pd.read_csv(data_path, encoding='GBK')
    print("Data loaded successfully using GBK encoding.")
except Exception as e:
    print("Failed to read the CSV file with GBK encoding:", e)
    try:
        hourly_prediction = pd.read_csv(data_path, encoding='utf-8-sig')
        print("Data loaded successfully using utf-8-sig encoding.")
    except Exception as e:
        print("Failed to read the CSV file with utf-8-sig encoding:", e)

# 设置列名和日期处理
hourly_prediction.columns = ['分拣中心', '日期', '班次', '预测货量']
hourly_prediction['日期'] = pd.to_datetime(hourly_prediction['日期'])

# 仅对分拣中心SC60的数据进行操作
sc60_data = hourly_prediction[hourly_prediction['分拣中心'] == 'SC60']

def solve_optimization(center, shifts_data):
    # 创建优化问题
    mylp = pp.LpProblem("Staffing Optimization", pp.LpMinimize)

    # 获取日期范围和班次
    dates = sorted(shifts_data['日期'].unique())
    shifts = shifts_data['班次'].unique()

    # 定义决策变量
    x = {(shift, date): pp.LpVariable(f"正式工_{shift}_{date.strftime('%Y%m%d')}", lowBound=0, cat="Integer")
         for date in dates for shift in shifts}
    y = {(shift, date): pp.LpVariable(f"临时工_{shift}_{date.strftime('%Y%m%d')}", lowBound=0, cat="Integer")
         for date in dates for shift in shifts}

    # 最小化总员工数
    mylp += pp.lpSum(x[shift, date] + y[shift, date] for shift in shifts for date in dates)

    # 满足需求的约束
    for date in dates:
        for shift in shifts:
            demand = shifts_data.loc[(shifts_data['日期'] == date) & (shifts_data['班次'] == shift), '预测货量'].item()
            mylp += 25 * x[shift, date] + 20 * y[shift, date] >= demand

    # 正式工总数约束
    mylp += pp.lpSum(x[shift, date] for shift in shifts for date in dates) <= 200

    # 求解问题
    mylp.solve(pp.PULP_CBC_CMD(msg=True))

    # 收集结果
    results = [{
        '分拣中心': center,
        '日期': date.strftime('%Y-%m-%d'),
        '班次': shift,
        '正式工人数': pp.value(x[shift, date]),
        '临时工人数': pp.value(y[shift, date])
    } for date in dates for shift in shifts]
    return results

# 对SC60分拣中心应用优化
results = []
for date in tqdm(sorted(sc60_data['日期'].unique()), desc="Optimizing SC60"):
    date_data = sc60_data[sc60_data['日期'] == date]
    results.extend(solve_optimization('SC60', date_data))

# 将结果列表保存
results_df = pd.DataFrame(results)
results_df.to_csv('optimized_staff_schedule_SC60.csv', index=False)
print("Optimization complete for SC60. Results saved to 'optimized_staff_schedule_SC60.csv'.")
